﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Dia2Lib;
using System.Threading;
using System.Runtime.InteropServices;
using Kokomo.Mirror.Ia32;

namespace Kokomo.Mirror.PeFormat.Symbols
{
	public class SymbolSource
	{
		class DiaSourceCreator
		{
			public void ThreadProc()
			{
				this.diaSource = new DiaSourceClass();

			}
			public DiaSourceClass diaSource;
		}

		public static SymbolSource LoadSymbolFile(string pdbFileName)
		{
			if (pdbFileName == null) throw new ArgumentNullException("pdbFileName");
			IDiaDataSource source = new DiaSourceClass();
			source.loadDataFromPdb(pdbFileName);
			IDiaSession session;
			source.openSession(out session);

			return new SymbolSource(pdbFileName, session);
		}


		private ComWrapper<IDiaSession> _sessionWrapper;
		internal IDiaSession Session
		{
			get { return this._sessionWrapper.Object; }
		}

		//internal IDiaSession Session;
		public string SourceFileName { get; private set; }

		internal SymbolSource(
			string sourceFileName,
			IDiaSession session
			)
		{
			this.SourceFileName = sourceFileName;

			this._sessionWrapper = new ComWrapper<IDiaSession>(session);
			this._tables = new Lazy<IEnumerable<Table>>(this.GetTables);
		}

		public ModuleSymbol GlobalScope
		{
			get { return (ModuleSymbol)Symbol.FromDiaSymbol(this.Session.globalScope, this); }
		}

		public SymbolicReference GetSymbolicReferenceForRva(
			uint rva
			)
		{
			Symbol symbol = this.FindSymbolByRelativeAddress(rva);
			if (symbol != null)
			{
				PublicSymbol publicSymbol = symbol as PublicSymbol;
				if (publicSymbol != null)
				{
					string name = symbol.Name;
					if (!string.IsNullOrEmpty(name))
					{
						int offset = (int)(rva - publicSymbol.RelativeVirtualAddress);
						SymbolicReference reference = new SymbolicReference(
							name,
							offset
							);

						return reference;
					}
				}
			}

			return null;
		}

		public TTable GetTable<TTable>()
			where TTable : Table
		{
			return this._tables.Value.OfType<TTable>().FirstOrDefault();
		}
		public IEnumerable<Symbol> FindSymbols(string name, bool exactMatch = true)
		{
			if (string.IsNullOrEmpty(name)) throw new ArgumentNullException("name");

			SymbolTable symbolTable = this.GetTable<SymbolTable>();
			if (symbolTable == null) return null;

			Func<string, string, bool> predicate =
				exactMatch
				? string.Equals
				: (Func<string, string, bool>)((x, y) => !string.IsNullOrEmpty(x) && x.Contains(y, StringComparison.InvariantCultureIgnoreCase))
				;
			IEnumerable<Symbol> symbol = symbolTable.Where(s => predicate(s.Name, name) || predicate(s.PrincipalName, name));
			return symbol;
		}
		public Symbol FindSymbol(string name, bool exactMatch = true)
		{
			return this.FindSymbols(name, exactMatch).FirstOrDefault();
		}

		public Symbol SymbolById(uint id)
		{
			IDiaSymbol diaSymbol = null;
			this.Session.symbolById(id, out diaSymbol);
			if (diaSymbol == null) return null;
			return Symbol.FromDiaSymbol(diaSymbol, this);
		}
		public Symbol FindSymbolByAddress(
			ulong virtualAddress,
			SymbolKind kind = SymbolKind.Null
			)
		{
			IDiaSymbol diaSymbol = null;
			this.Session.findSymbolByVA(virtualAddress, (SymTagEnum)kind, out diaSymbol);
			if (diaSymbol == null) return null;
			return Symbol.FromDiaSymbol(diaSymbol, this);
		}
		public Symbol FindSymbolByAddress(
			ulong virtualAddress,
			out int displacement,
			SymbolKind kind = SymbolKind.Null
			)
		{
			IDiaSymbol diaSymbol = null;
			this.Session.findSymbolByVAEx(virtualAddress, (SymTagEnum)kind, out diaSymbol, out displacement);
			if (diaSymbol == null) return null;
			return Symbol.FromDiaSymbol(diaSymbol, this);
		}
		public Symbol FindSymbolByRelativeAddress(
			uint virtualAddress,
			SymbolKind kind = SymbolKind.Null
			)
		{
			IDiaSymbol diaSymbol = null;
			this.Session.findSymbolByRVA(virtualAddress, (SymTagEnum)kind, out diaSymbol);
			if (diaSymbol == null) return null;
			return Symbol.FromDiaSymbol(diaSymbol, this);
		}
		public Symbol FindSymbolByRelativeAddress(
			uint virtualAddress,
			out int displacement,
			SymbolKind kind = SymbolKind.Null
			)
		{
			IDiaSymbol diaSymbol = null;
			this.Session.findSymbolByRVAEx(virtualAddress, (SymTagEnum)kind, out diaSymbol, out displacement);
			if (diaSymbol == null) return null;
			return Symbol.FromDiaSymbol(diaSymbol, this);
		}

		public IEnumerable<SourceFile> GetSourceFiles()
		{
			IDiaEnumSourceFiles sourceFiles = null;
			this.Session.findFile(null, null, 0, out sourceFiles);
			if (sourceFiles == null) yield break;

			while (true)
			{
				uint fetched;
				IDiaSourceFile sourceFile;
				sourceFiles.Next(1, out sourceFile, out fetched);
				if (fetched != 1) yield break;

				yield return new SourceFile(sourceFile);
			}
		}
		public InjectedSourceFile GetInjectedSourceFile(string sourceFileName)
		{
			if (sourceFileName == null) throw new ArgumentNullException("sourceFileName");

			IDiaEnumInjectedSources injectedSources = null;
			this.Session.findInjectedSource(sourceFileName, out injectedSources);
			if (injectedSources == null) return null;

			uint fetched;
			IDiaInjectedSource injectedSource;
			injectedSources.Next(1, out injectedSource, out fetched);
			if (fetched != 1) return null;
			if (injectedSource == null) return null;

			return new InjectedSourceFile(injectedSource);
		}

		private Lazy<IEnumerable<Table>> _tables;
		public IEnumerable<Table> Tables
		{
			get { return this._tables.Value; }
		}
		private IEnumerable<Table> GetTables()
		{
			return new TableEnumerable(this);
		}



		//void findFile(IDiaSymbol pCompiland, string name, uint compareFlags, out IDiaEnumSourceFiles ppResult);
		//void findFileById(uint uniqueId, out IDiaSourceFile ppResult);
		//void findInjectedSource(string srcFile, out IDiaEnumInjectedSources ppResult);
		//void findLines(IDiaSymbol compiland, IDiaSourceFile file, out IDiaEnumLineNumbers ppResult);
		//void findLinesByAddr(uint seg, uint offset, uint length, out IDiaEnumLineNumbers ppResult);
		//void findLinesByLinenum(IDiaSymbol compiland, IDiaSourceFile file, uint linenum, uint column, out IDiaEnumLineNumbers ppResult);
		//void findLinesByRVA(uint rva, uint length, out IDiaEnumLineNumbers ppResult);
		//void findLinesByVA(ulong va, uint length, out IDiaEnumLineNumbers ppResult);
		//void findSymbolByAddr(uint isect, uint offset, SymTagEnum symTag, out IDiaSymbol ppSymbol);
		//void findSymbolByRVA(uint rva, SymTagEnum symTag, out IDiaSymbol ppSymbol);
		//void findSymbolByRVAEx(uint rva, SymTagEnum symTag, out IDiaSymbol ppSymbol, out int displacement);
		//void findSymbolByToken(uint token, SymTagEnum symTag, out IDiaSymbol ppSymbol);
		//void getEnumDebugStreams(out IDiaEnumDebugStreams ppEnumDebugStreams);
		//void getEnumTables(out IDiaEnumTables ppEnumTables);
		//void getSymbolsByAddr(out IDiaEnumSymbolsByAddr ppEnumbyAddr);
		//void symsAreEquiv(IDiaSymbol symbolA, IDiaSymbol symbolB);
	}
}
